from PySide6.QtWidgets import QTableView, QHeaderView
from PySide6.QtCore import Qt, QAbstractTableModel, QModelIndex
import pandas as pd

class PandasTableModel(QAbstractTableModel):
    """Pandas数据表模型"""
    
    def __init__(self, data: pd.DataFrame = None):
        super().__init__()
        self._data = data if data is not None else pd.DataFrame()
    
    def set_data(self, data: pd.DataFrame):
        """设置新数据"""
        self.beginResetModel()
        self._data = data
        self.endResetModel()
    
    def rowCount(self, parent=QModelIndex()) -> int:
        """行数"""
        return len(self._data)
    
    def columnCount(self, parent=QModelIndex()) -> int:
        """列数"""
        return len(self._data.columns)
    
    def data(self, index: QModelIndex, role: int = Qt.DisplayRole):
        """获取数据"""
        if not index.isValid():
            return None
        
        if role == Qt.DisplayRole or role == Qt.EditRole:
            value = self._data.iloc[index.row(), index.column()]
            return str(value) if not pd.isna(value) else ""
        
        return None
    
    def headerData(self, section: int, orientation: Qt.Orientation, role: int = Qt.DisplayRole):
        """获取表头数据"""
        if role != Qt.DisplayRole:
            return None
        
        if orientation == Qt.Horizontal:
            return str(self._data.columns[section])
        else:
            return str(self._data.index[section])
    
    def flags(self, index: QModelIndex) -> Qt.ItemFlags:
        """项标志"""
        return Qt.ItemIsEnabled | Qt.ItemIsSelectable | Qt.ItemIsEditable

class DataTableView(QTableView):
    """数据表格视图"""
    
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setModel(PandasTableModel())
        self._setup_ui()
    
    def _setup_ui(self):
        """设置UI"""
        # 启用交替行颜色
        self.setAlternatingRowColors(True)
        
        # 设置选择行为
        self.setSelectionBehavior(QTableView.SelectRows)
        self.setSelectionMode(QTableView.SingleSelection)
        
        # 调整列宽适应内容
        self.horizontalHeader().setSectionResizeMode(QHeaderView.Interactive)
        self.horizontalHeader().setStretchLastSection(True)
        
        # 启用排序
        self.setSortingEnabled(True)
    
    def set_data(self, data: pd.DataFrame):
        """设置数据"""
        self.model().set_data(data)
        
        # 调整列宽
        self.resizeColumnsToContents()
        
        # 如果列数太多，改为可滚动
        if len(data.columns) > 10:
            self.horizontalHeader().setSectionResizeMode(QHeaderView.Interactive)
        else:
            self.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)